{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a945ea64",
   "metadata": {},
   "source": [
    "# LightNER Based NER Tutorial\n",
    "> Tutorial author: 李磊（<leili21@zju.edu.cn>）\n",
    "\n",
    "In this tutorial, we use a generative method [LightNER](https://arxiv.org/abs/2109.00720) to recognize named entities. We hope this tutorial can help you understand the process of named entity recognition.\n",
    "\n",
    "This tutorial uses `Python3`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62dcbdb5",
   "metadata": {},
   "source": [
    "## NER\n",
    "**Named-entity recognition** (also known as named entity identification, entity chunking, and entity extraction) is a subtask of information extraction that seeks to locate and classify named entities mentioned in unstructured text into pre-defined categories such as person names, organizations, locations, medical codes, time expressions, quantities, monetary values, percentages, etc.\n",
    "\n",
    "## Few-shot NER\n",
    "In few-shot NER, you only have a few training instances. For example, each entity type provides k samples, which we call k-shot NER. In this demo, we try two methods. The first one is loading the model trained from rich resources then training in few-shot NER datasets. The seceond one is directly fine-tuning pretrained models in few-shot NER datasets."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e6afae3",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "In rich resource scenarios, we use dataset [**CoNLL-2003**](https://www.clips.uantwerpen.be/conll2003/ner/). It is a dataset for NER, concentrating on four types of named entities related to persons, locations, organizations, and names of miscellaneous entities.\n",
    "\n",
    "|    Word    | Named entity tag |\n",
    "| :--------: | :--------------: |\n",
    "|  Pakistan  |      B-LOC       |\n",
    "|     ,      |        O         |\n",
    "|    who     |        O         |\n",
    "|   arrive   |        O         |\n",
    "|     in     |        O         |\n",
    "| Australia  |      B-LOC       |\n",
    "|   later    |        O         |\n",
    "|    this    |        O         |\n",
    "|   month    |        O         |\n",
    "|     ,      |        O         |\n",
    "|    are     |        O         |\n",
    "|    the     |        O         |\n",
    "|   other    |        O         |\n",
    "|    team    |        O         |\n",
    "| competing  |        O         |\n",
    "|     in     |        O         |\n",
    "|    the     |        O         |\n",
    "|   World    |      B-MISC      |\n",
    "|   Series   |      I-MISC      |\n",
    "| tournament |        O         |\n",
    "|     .      |        O         |\n",
    "\n",
    "- train.txt: It contains 14,987 sentences\n",
    "- valid.txt: It contains 3,466 sentences\n",
    "- test.txt: It contains 3,684 sentences\n",
    "\n",
    "In low resource scenarios, we use datasets [ATIS](https://github.com/yvchen/JointSLU/tree/master/data), [MIT Restaurant](https://groups.csail.mit.edu/sls/downloads/) and [MIT Movie](https://groups.csail.mit.edu/sls/downloads/), each dataset is preprocessed into conll format and sampled k-shot data for training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b62fef7",
   "metadata": {},
   "source": [
    "## LightNER\n",
    "LightNER: A Lightweight Generative Framework with Prompt-guided Attention\n",
    "for Low-resource NER (LightNER) is novel lightweight generative framework with prompt-guided attention for low-resource NER.\n",
    "\n",
    "![LightNER](img/LightNER.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5946b060",
   "metadata": {},
   "source": [
    "## Prepare the runtime environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0c311a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install deepke\n",
    "!wget 120.27.214.45/Data/ner/few_shot/data.tar.gz\n",
    "!tar -xzvf data.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46f37a94",
   "metadata": {},
   "source": [
    "## Import modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b0cd7f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]='0'\n",
    "import logging\n",
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "from mapping_type import mit_movie_mapping, mit_restaurant_mapping, atis_mapping\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
    "from tensorboardX import SummaryWriter\n",
    "writer = SummaryWriter(log_dir='logs')\n",
    "\n",
    "logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',\n",
    "                    datefmt = '%m/%d/%Y %H:%M:%S',\n",
    "                    level = logging.INFO)\n",
    "logger = logging.getLogger(__name__)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "213c1403",
   "metadata": {},
   "source": [
    "## Processer and Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dea6a7ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "from itertools import chain\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from transformers import BartTokenizer\n",
    "\n",
    "import logging\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "# load file and process bio\n",
    "class ConllNERProcessor(object):\n",
    "    def __init__(self, data_path, mapping, bart_name, learn_weights) -> None:\n",
    "        self.data_path = data_path\n",
    "        self.tokenizer = BartTokenizer.from_pretrained(bart_name)\n",
    "        self.mapping = mapping  # 记录的是原始tag与转换后的tag的str的匹配关系\n",
    "        self.original_token_nums = self.tokenizer.vocab_size\n",
    "        self.learn_weights = learn_weights\n",
    "        self._add_tags_to_tokens()\n",
    "\n",
    "    def load_from_file(self, mode='train'):\n",
    "        \"\"\"load conll ner from file\n",
    "\n",
    "        Args:\n",
    "            mode (str, optional): train/test/dev. Defaults to 'train'.\n",
    "        Return:\n",
    "            outputs (dict)\n",
    "            raw_words: ['EU', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.']\n",
    "            raw_targets: ['B-ORG', 'O', 'B-MISC', 'O', 'O', 'O', 'B-MISC', 'O', 'O']\n",
    "            entities: [['EU'], ['German'], ['British']]\n",
    "            entity_tags: ['org', 'misc', 'misc']\n",
    "            entity_spans: [[0, 1], [2, 3], [6, 7]]\n",
    "        \"\"\"\n",
    "        load_file = self.data_path[mode]\n",
    "        logger.info(\"Loading data from {}\".format(load_file))\n",
    "\n",
    "        # extract bio\n",
    "        split_c = '\\t' if 'conll' in load_file  else ' '\n",
    "        outputs = {'raw_words':[], 'raw_targets':[], 'entities':[], 'entity_tags':[], 'entity_spans':[]}\n",
    "        with open(load_file, \"r\", encoding=\"utf-8\") as f:\n",
    "            lines = f.readlines()\n",
    "            raw_words, raw_targets = [], []\n",
    "            raw_word, raw_target = [], []\n",
    "            for line in lines:\n",
    "                if line != \"\\n\":\n",
    "                    raw_word.append(line.split(split_c)[0]) # conll \"\\t\" else \" \"\n",
    "                    raw_target.append(line.split(split_c)[1][:-1])\n",
    "                else:\n",
    "                    raw_words.append(raw_word)\n",
    "                    raw_targets.append(raw_target)\n",
    "                    raw_word, raw_target = [], []\n",
    "\n",
    "        for words, targets in zip(raw_words, raw_targets):\n",
    "            entities, entity_tags, entity_spans = [], [], []\n",
    "            start, end, start_flag = 0, 0, False\n",
    "            for idx, tag in enumerate(targets):\n",
    "                if tag.startswith('B-'):    # 一个实体开头 另一个实体（I-）结束\n",
    "                    end = idx\n",
    "                    if start_flag:  # 另一个实体以I-结束，紧接着当前实体B-出现\n",
    "                        entities.append(words[start:end])\n",
    "                        entity_tags.append(targets[start][2:].lower())\n",
    "                        entity_spans.append([start, end])\n",
    "                        start_flag = False\n",
    "                    start = idx\n",
    "                    start_flag = True\n",
    "                elif tag.startswith('I-'):  # 实体中间，不是开头也不是结束，end+1即可\n",
    "                    end = idx\n",
    "                elif tag.startswith('O'):  # 无实体，可能是上一个实体的结束\n",
    "                    end = idx\n",
    "                    if start_flag:  # 上一个实体结束\n",
    "                        entities.append(words[start:end])\n",
    "                        entity_tags.append(targets[start][2:].lower())\n",
    "                        entity_spans.append([start, end])\n",
    "                        start_flag = False\n",
    "            if start_flag:  # 句子以实体I-结束，未被添加\n",
    "                entities.append(words[start:end+1])\n",
    "                entity_tags.append(targets[start][2:].lower())\n",
    "                entity_spans.append([start, end+1])\n",
    "                start_flag = False\n",
    "    \n",
    "            if len(entities) != 0:\n",
    "                outputs['raw_words'].append(words)\n",
    "                outputs['raw_targets'].append(targets)\n",
    "                outputs['entities'].append(entities)\n",
    "                outputs['entity_tags'].append(entity_tags)\n",
    "                outputs['entity_spans'].append(entity_spans)\n",
    "        return outputs\n",
    "\n",
    "    def process(self, data_dict):\n",
    "        target_shift = len(self.mapping) + 2 \n",
    "        def prepare_target(item):\n",
    "            raw_word = item['raw_word']\n",
    "            word_bpes = [[self.tokenizer.bos_token_id]] \n",
    "            first = [] \n",
    "            cur_bpe_len = 1\n",
    "            for word in raw_word:\n",
    "                bpes = self.tokenizer.tokenize(word, add_prefix_space=True)\n",
    "                bpes = self.tokenizer.convert_tokens_to_ids(bpes)\n",
    "                first.append(cur_bpe_len)\n",
    "                cur_bpe_len += len(bpes)\n",
    "                word_bpes.append(bpes)\n",
    "            assert first[-1] + len(bpes) == sum(map(len, word_bpes))\n",
    "            word_bpes.append([self.tokenizer.eos_token_id])\n",
    "            assert len(first) == len(raw_word) == len(word_bpes) - 2\n",
    "\n",
    "            lens = list(map(len, word_bpes)) \n",
    "            cum_lens = np.cumsum(lens).tolist()   \n",
    "\n",
    "            entity_spans = item['entity_span']  # [(s1, e1, s2, e2), ()]\n",
    "            entity_tags = item['entity_tag']  # [tag1, tag2...]\n",
    "            entities = item['entity']  # [[ent1, ent2,], [ent1, ent2]]\n",
    "            target = [0]\n",
    "            pairs = []\n",
    "\n",
    "            first = list(range(cum_lens[-1]))\n",
    "\n",
    "            assert len(entity_spans) == len(entity_tags)                #\n",
    "            for idx, (entity, tag) in enumerate(zip(entity_spans, entity_tags)):\n",
    "                cur_pair = []\n",
    "                num_ent = len(entity) // 2\n",
    "                for i in range(num_ent):\n",
    "                    start = entity[2 * i]\n",
    "                    end = entity[2 * i + 1]\n",
    "                    cur_pair_ = []\n",
    "                    cur_pair_.extend([cum_lens[k] for k in list(range(start, end))])\n",
    "                    cur_pair.extend([p + target_shift for p in cur_pair_])\n",
    "                for _, (j, word_idx) in enumerate(zip((cur_pair[0], cur_pair[-1]), (0, -1))):\n",
    "                    j = j - target_shift\n",
    "                assert all([cur_pair[i] < cum_lens[-1] + target_shift for i in range(len(cur_pair))])\n",
    "\n",
    "                cur_pair.append(self.mapping2targetid[tag] + 2) \n",
    "                pairs.append([p for p in cur_pair])\n",
    "            target.extend(list(chain(*pairs)))\n",
    "            target.append(1) \n",
    "\n",
    "            word_bpes = list(chain(*word_bpes))\n",
    "            assert len(word_bpes)<500\n",
    "\n",
    "            dict  = {'tgt_tokens': target, 'target_span': pairs, 'src_tokens': word_bpes,\n",
    "                    'first': first, 'src_seq_len':len(word_bpes), 'tgt_seq_len':len(target)}\n",
    "            return dict\n",
    "        \n",
    "        logger.info(\"Process data...\")\n",
    "        for raw_word, raw_target, entity, entity_tag, entity_span in tqdm(zip(data_dict['raw_words'], data_dict['raw_targets'], data_dict['entities'], \n",
    "                                                                                data_dict['entity_tags'], data_dict['entity_spans']), total=len(data_dict['raw_words']), desc='Processing'):\n",
    "            item_dict = prepare_target({'raw_word': raw_word, 'raw_target':raw_target, 'entity': entity, 'entity_tag': entity_tag, 'entity_span': entity_span})\n",
    "            # add item_dict to data_dict\n",
    "            for key, value in item_dict.items():\n",
    "                if key in data_dict:\n",
    "                    data_dict[key].append(value)\n",
    "                else:\n",
    "                    data_dict[key] = [value]\n",
    "        return data_dict\n",
    "\n",
    "    def _add_tags_to_tokens(self):\n",
    "        mapping = self.mapping\n",
    "        if self.learn_weights:  # add extra tokens to huggingface tokenizer\n",
    "            self.mapping2id = {} \n",
    "            self.mapping2targetid = {} \n",
    "            for key, value in self.mapping.items():\n",
    "                key_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(value[2:-2], add_prefix_space=True))\n",
    "                self.mapping2id[value] = key_id  # may be list\n",
    "                self.mapping2targetid[key] = len(self.mapping2targetid)\n",
    "        else:\n",
    "            tokens_to_add = sorted(list(mapping.values()), key=lambda x: len(x), reverse=True)  # \n",
    "            unique_no_split_tokens = self.tokenizer.unique_no_split_tokens                      # no split\n",
    "            sorted_add_tokens = sorted(list(tokens_to_add), key=lambda x: len(x), reverse=True)\n",
    "            for tok in sorted_add_tokens:\n",
    "                assert self.tokenizer.convert_tokens_to_ids([tok])[0] == self.tokenizer.unk_token_id    # \n",
    "            self.tokenizer.unique_no_split_tokens = unique_no_split_tokens + sorted_add_tokens          # add to no_split_tokens\n",
    "            self.tokenizer.add_tokens(sorted_add_tokens)\n",
    "            self.mapping2id = {}  # tag to id\n",
    "            self.mapping2targetid = {}  # tag to number\n",
    "\n",
    "            for key, value in self.mapping.items():\n",
    "                key_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(value))\n",
    "                assert len(key_id) == 1, value\n",
    "                assert key_id[0] >= self.original_token_nums\n",
    "                self.mapping2id[value] = key_id[0]  #\n",
    "                self.mapping2targetid[key] = len(self.mapping2targetid)\n",
    "        \n",
    "\n",
    "class ConllNERDataset(Dataset):\n",
    "    def __init__(self, data_processor, mode='train') -> None:\n",
    "        self.data_processor = data_processor\n",
    "        self.data_dict = data_processor.load_from_file(mode=mode)\n",
    "        self.complet_data = data_processor.process(self.data_dict)\n",
    "        self.mode = mode\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.complet_data['src_tokens'])\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        if self.mode == 'test':\n",
    "            return torch.tensor(self.complet_data['src_tokens'][index]), torch.tensor(self.complet_data['src_seq_len'][index]), \\\n",
    "                    torch.tensor(self.complet_data['first'][index]), self.complet_data['raw_words'][index]\n",
    "\n",
    "        return torch.tensor(self.complet_data['src_tokens'][index]), torch.tensor(self.complet_data['tgt_tokens'][index]), \\\n",
    "                    torch.tensor(self.complet_data['src_seq_len'][index]), torch.tensor(self.complet_data['tgt_seq_len'][index]), \\\n",
    "                    torch.tensor(self.complet_data['first'][index]), self.complet_data['target_span'][index]\n",
    "\n",
    "\n",
    "    def collate_fn(self, batch):\n",
    "        src_tokens, src_seq_len, first  = [], [], []\n",
    "        tgt_tokens, tgt_seq_len, target_span = [], [], []\n",
    "        if self.mode == \"test\":\n",
    "            raw_words = []\n",
    "            for tup in batch:\n",
    "                src_tokens.append(tup[0])\n",
    "                src_seq_len.append(tup[1])\n",
    "                first.append(tup[2])\n",
    "                raw_words.append(tup[3])\n",
    "            src_tokens = pad_sequence(src_tokens, batch_first=True, padding_value=self.data_processor.tokenizer.pad_token_id)\n",
    "            first = pad_sequence(first, batch_first=True, padding_value=0)\n",
    "            return src_tokens, torch.stack(src_seq_len, 0), first, raw_words\n",
    "\n",
    "        for tup in batch:\n",
    "            src_tokens.append(tup[0])\n",
    "            tgt_tokens.append(tup[1])\n",
    "            src_seq_len.append(tup[2])\n",
    "            tgt_seq_len.append(tup[3])\n",
    "            first.append(tup[4])\n",
    "            target_span.append(tup[5])\n",
    "        src_tokens = pad_sequence(src_tokens, batch_first=True, padding_value=self.data_processor.tokenizer.pad_token_id)\n",
    "        tgt_tokens = pad_sequence(tgt_tokens, batch_first=True, padding_value=1)\n",
    "        first = pad_sequence(first, batch_first=True, padding_value=0)\n",
    "        return src_tokens, tgt_tokens, torch.stack(src_seq_len, 0), torch.stack(tgt_seq_len, 0), first, target_span\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99815b70",
   "metadata": {},
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d33345b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import random\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "from transformers import BartModel, BartTokenizer\n",
    "import os\n",
    "\n",
    "def avg_token_embeddings(tokenizer: BartTokenizer, bart_model: BartModel, bart_name, num_tokens):\n",
    "    \"\"\"when initial added tokens, use their averge token emebddings\n",
    "\n",
    "    Args:\n",
    "        tokenizer (BartTokenizer): [description]\n",
    "        bart_model (BartModel): [description]\n",
    "        bart_name ([type]): [description]\n",
    "        num_tokens ([type]): [description]\n",
    "\n",
    "    Raises:\n",
    "        RuntimeError: [description]\n",
    "\n",
    "    Returns:\n",
    "        [type]: [description]\n",
    "    \"\"\"\n",
    "    _tokenizer = BartTokenizer.from_pretrained(bart_name)\n",
    "    for token in tokenizer.unique_no_split_tokens:\n",
    "        if token[:2] == '<<':  # 特殊字符\n",
    "            index = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(token))\n",
    "            if len(index)>1:\n",
    "                raise RuntimeError(f\"{token} wrong split\")\n",
    "            else:\n",
    "                index = index[0]\n",
    "            assert index>=num_tokens, (index, num_tokens, token)\n",
    "            indexes = _tokenizer.convert_tokens_to_ids(_tokenizer.tokenize(token[2:-2]))\n",
    "            embed = bart_model.encoder.embed_tokens.weight.data[indexes[0]]\n",
    "            for i in indexes[1:]:\n",
    "                embed += bart_model.decoder.embed_tokens.weight.data[i]\n",
    "            embed /= len(indexes)\n",
    "            bart_model.decoder.embed_tokens.weight.data[index] = embed\n",
    "    return bart_model\n",
    "\n",
    "\n",
    "def seq_to_mask(seq_len, max_len):\n",
    "    \"\"\"[get attention mask with sequence length]\n",
    "\n",
    "    Args:\n",
    "        seq_len ([torch.tensor]): [shape: bsz, each sequence length in a batch]\n",
    "    \"\"\"\n",
    "    max_len = int(max_len) if max_len else seq_len.max().long()\n",
    "    cast_seq = torch.arange(max_len).expand(seq_len.size(0), -1).to(seq_len)\n",
    "    mask = cast_seq.lt(seq_len.unsqueeze(1))\n",
    "    return mask\n",
    "\n",
    "def get_loss(tgt_tokens, tgt_seq_len, pred):\n",
    "    \"\"\"\n",
    "\n",
    "    :param tgt_tokens: bsz x max_len, 包含了的[sos, token, eos]\n",
    "    :param pred: bsz x max_len-1 x vocab_size\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    tgt_seq_len = tgt_seq_len - 1\n",
    "    mask = seq_to_mask(tgt_seq_len, max_len=tgt_tokens.size(1) - 1).eq(0)\n",
    "    tgt_tokens = tgt_tokens[:, 1:].masked_fill(mask, -100)\n",
    "    loss = F.cross_entropy(target=tgt_tokens, input=pred.transpose(1, 2))\n",
    "    return loss\n",
    "\n",
    "def get_model_device(model):\n",
    "    assert isinstance(model, nn.Module)\n",
    "\n",
    "    parameters = list(model.parameters())\n",
    "    if len(parameters) == 0:\n",
    "        return None\n",
    "    else:\n",
    "        return parameters[0].device\n",
    "\n",
    "def set_seed(seed=2021):\n",
    "    \"\"\"sets random seed\"\"\"\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    np.random.seed(seed)\n",
    "    random.seed(seed)\n",
    "\n",
    "\n",
    "def convert_preds_to_outputs(preds, raw_words, mapping, tokenizer):\n",
    "    \"\"\"convet model predicitons to BIO outputs\n",
    "\n",
    "    Args:\n",
    "        preds ([torch.Tensor]): [prompt model predictions, (bsz x seq_len x labels)]\n",
    "        raw_words ([List]): [source raw words]\n",
    "        mapping ([dict]): [map entity labels to <<>>]\n",
    "        tokenizer : [BartTokenizer]\n",
    "\n",
    "    Returns:\n",
    "        [outputs (List)]: [each item length equal to raw_words, BIO format.]\n",
    "    \"\"\"\n",
    "    id2label = list(mapping.keys())\n",
    "    pred_eos_index = preds.flip(dims=[1]).eq(1).cumsum(dim=1).long()\n",
    "    preds = preds[:, 1:]\n",
    "    pred_seq_len = pred_eos_index.flip(dims=[1]).eq(pred_eos_index[:, -1:]).sum(dim=1)  # bsz\n",
    "    pred_seq_len = (pred_seq_len - 2).tolist()\n",
    "\n",
    "    word_start_index = len(mapping) + 2\n",
    "    outputs = []\n",
    "    for i, pred_item in enumerate(preds.tolist()):\n",
    "        pred_item = pred_item[:pred_seq_len[i]] # single sentence prediction\n",
    "        pairs, cur_pair = [], []\n",
    "        if len(pred_item):  # this sentence prediciton= is not null\n",
    "            for idx in pred_item:\n",
    "                if idx < word_start_index:  # is entity\n",
    "                    if len(cur_pair) > 0:\n",
    "                        # assert word[i] < word[i+1]\n",
    "                        if all([cur_pair[i] < cur_pair[i + 1] for i in range(len(cur_pair) - 1)]):\n",
    "                            pairs.append(tuple(cur_pair + [idx]))   # add valid words and current entity id\n",
    "                    cur_pair = []   # clear word pairs\n",
    "                else:   # is word\n",
    "                    cur_pair.append(idx)    # add word id to word pairs\n",
    "        raw_words_item = raw_words[i]\n",
    "        cum_lens = [1]\n",
    "        start_idx = 1\n",
    "        for word in raw_words_item:\n",
    "            start_idx += len(tokenizer.tokenize(word, add_prefix_space=True))\n",
    "            cum_lens.append(start_idx)\n",
    "        cum_lens.append(start_idx+1)\n",
    "        output = ['O' for _ in range(len(raw_words_item))]\n",
    "        # pairs: List[(word id, ... , entity id), (...), ...]\n",
    "        for pair in pairs:  # (word id, ... , entity id)\n",
    "            entity = pair[-1]\n",
    "            words = []\n",
    "            for word in pair[:-1]:\n",
    "                if word-word_start_index in cum_lens:\n",
    "                    words.append(cum_lens.index(word-word_start_index)) \n",
    "            if len(words) == 0: continue\n",
    "            start_idx = words[0]\n",
    "            end_idx = words[-1]\n",
    "            output[start_idx] = f'B-{id2label[entity-2]}'\n",
    "            for _ in range(start_idx+1, end_idx+1):\n",
    "                output[_] = f'I-{id2label[entity-2]}'\n",
    "        outputs.append(output)\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def write_predictions(path, texts, labels):\n",
    "    \"\"\"[write model predictions to path (conll format)]\n",
    "\n",
    "    Args:\n",
    "        path ([str]): [save path]\n",
    "        texts ([List]): [raw texts]\n",
    "        labels ([List]): [predict labels]\n",
    "    \"\"\"\n",
    "    print(len(texts), len(labels))\n",
    "    assert len(texts) == len(labels)\n",
    "    if not os.path.exists(path):\n",
    "        os.system(r\"touch {}\".format(path))\n",
    "    \n",
    "    with open(path, \"w\", encoding=\"utf-8\") as f:\n",
    "        f.writelines(\"-DOCSTART-\tO\\n\\n\")\n",
    "        for i in range(len(texts)):\n",
    "            for j in range(len(texts[i])):\n",
    "                f.writelines(\"{}\\t{}\\n\".format(texts[i][j], labels[i][j]))\n",
    "            f.writelines(\"\\n\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4846198a",
   "metadata": {},
   "source": [
    "## LightNER model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93e2c547",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from transformers.configuration_bart import BartConfig\n",
    "from modeling_bart import BartModel, _prepare_bart_decoder_inputs\n",
    "\n",
    "from functools import partial\n",
    "from typing import Union\n",
    "\n",
    "\n",
    "class PromptBartEncoder(nn.Module):\n",
    "    def __init__(self, encoder):\n",
    "        super(PromptBartEncoder, self).__init__()\n",
    "        self.bart_encoder = encoder\n",
    "    \n",
    "    def forward(self, src_tokens, attention_mask=None, past_key_values=None):\n",
    "        encoder_dicts = self.bart_encoder(input_ids=src_tokens, attention_mask=attention_mask, past_key_values=past_key_values, return_dict=True, output_hidden_states=True)\n",
    "        return encoder_dicts.last_hidden_state, encoder_dicts.hidden_states\n",
    "        \n",
    "class PromptBartDecoder(nn.Module):\n",
    "    def __init__(self, decoder, pad_token_id, label_ids, use_prompt=False, prompt_len=10, learn_weights=False):\n",
    "        super(PromptBartDecoder, self).__init__()\n",
    "        self.bart_decoder = decoder\n",
    "        self.pad_token_id = pad_token_id\n",
    "        self.use_prompt = use_prompt\n",
    "        self.prompt_len = prompt_len\n",
    "        self.learn_weights = learn_weights\n",
    "        self.label_ids = label_ids\n",
    "\n",
    "        print(label_ids)\n",
    "        if self.learn_weights:   # set learnable averge weights\n",
    "            self.averge_weights = nn.ParameterList(parameters=None)\n",
    "            for id in label_ids:\n",
    "                if len(id) > 1:\n",
    "                    self.averge_weights.append(nn.Parameter(torch.FloatTensor(len(id)).uniform_(1.0, 2.5)))\n",
    "            print(self.averge_weights)\n",
    "            mapping = [0, 2]\n",
    "            for id in label_ids:\n",
    "                mapping += id[:1]\n",
    "            mapping = torch.LongTensor(mapping)\n",
    "        else:\n",
    "            mapping = torch.LongTensor([0, 2]+label_ids)\n",
    "            self.label_start_id = min(label_ids)\n",
    "            self.label_end_id = max(label_ids)+1\n",
    "\n",
    "        self.register_buffer('mapping', mapping)\n",
    "        self.src_start_index = len(mapping)\n",
    "\n",
    "        hidden_size = decoder.embed_tokens.weight.size(1)\n",
    "        self.bart_mlp = nn.Sequential(nn.Linear(hidden_size, hidden_size),\n",
    "                                    nn.Dropout(0.3),\n",
    "                                    nn.ReLU(),\n",
    "                                    nn.Linear(hidden_size, hidden_size))\n",
    "        self.dropout_layer = nn.Dropout(0.3)\n",
    "    \n",
    "    def forward(self, tgt_tokens, prompt_state):\n",
    "        cumsum = tgt_tokens.eq(1).flip(dims=[1]).cumsum(dim=-1)\n",
    "        tgt_pad_mask = cumsum.flip(dims=[-1]).ne(cumsum[:, -1:])\n",
    "\n",
    "        encoder_outputs = prompt_state.encoder_output  # last_hidden_state\n",
    "        attention_mask = prompt_state.encoder_mask   # attention_mask\n",
    "        first = prompt_state.first\n",
    "        src_tokens = prompt_state.src_tokens\n",
    "        past_key_values = prompt_state.past_key_values\n",
    "\n",
    "        # mapping target tokens\n",
    "        mapping_token_mask = tgt_tokens.lt(self.src_start_index) \n",
    "        mapped_tokens = tgt_tokens.masked_fill(tgt_tokens.ge(self.src_start_index), 0)\n",
    "        tag_mapped_tokens = self.mapping[mapped_tokens]\n",
    "\n",
    "        src_tokens_index = tgt_tokens - self.src_start_index # bsz x num_src_token\n",
    "        src_tokens_index = src_tokens_index.masked_fill(src_tokens_index.lt(0), 0)\n",
    "        if first is not None:\n",
    "            src_tokens = src_tokens.gather(index=first, dim=1)\n",
    "        word_mapped_tokens = src_tokens.gather(index=src_tokens_index, dim=1)\n",
    "\n",
    "        tokens = torch.where(mapping_token_mask, tag_mapped_tokens, word_mapped_tokens)  # bsz x max_len\n",
    "        tokens = tokens.masked_fill(tgt_pad_mask, self.pad_token_id)\n",
    "\n",
    "        decoder_input_ids, _, causal_mask = _prepare_bart_decoder_inputs(\n",
    "                self.pad_token_id, \n",
    "                tokens,\n",
    "                decoder_input_ids=None,\n",
    "                decoder_padding_mask=None,\n",
    "                causal_mask_dtype=self.bart_decoder.embed_tokens.weight.dtype\n",
    "        )\n",
    "\n",
    "        if self.use_prompt:\n",
    "            assert past_key_values is not None\n",
    "            _, _, seqlen, _ = past_key_values[0]['self']['prev_value'].shape\n",
    "            tgt_len = decoder_input_ids.size(1)\n",
    "            temp_mask = torch.zeros(tgt_len, seqlen).to(causal_mask.device) #tgtlen, preseqlen\n",
    "            causal_mask = torch.cat([temp_mask, causal_mask],dim=1) #tgtlen, preseqlen+tgtlen\n",
    "\n",
    "        if self.training:\n",
    "            tokens = tokens[:, :-1]\n",
    "            decoder_pad_mask = tokens.eq(self.pad_token_id) \n",
    "            dict = self.bart_decoder(input_ids=tokens,\n",
    "                                encoder_hidden_states=encoder_outputs,  # last_hidden_state\n",
    "                                encoder_padding_mask=attention_mask,  # attention_mask\n",
    "                                decoder_padding_mask=decoder_pad_mask,\n",
    "                                decoder_causal_mask=causal_mask[:tokens.size(1), :self.prompt_len+tokens.size(1)],\n",
    "                                output_hidden_states=True,\n",
    "                                past_key_values=past_key_values,\n",
    "                                return_dict=True)\n",
    "        else:\n",
    "            past_key_values = prompt_state.past_key_values\n",
    "            dict = self.bart_decoder(input_ids=tokens,\n",
    "                                encoder_hidden_states=encoder_outputs,\n",
    "                                encoder_padding_mask=attention_mask,\n",
    "                                decoder_padding_mask=None,\n",
    "                                decoder_causal_mask=None,\n",
    "                                past_key_values=past_key_values,\n",
    "                                use_cache=True,\n",
    "                                return_dict=True)\n",
    "        hidden_state = dict.last_hidden_state  # bsz x max_len x hidden_size\n",
    "        hidden_state = self.dropout_layer(hidden_state)\n",
    "        if not self.training:\n",
    "            prompt_state.past_key_values = dict.past_key_values\n",
    "\n",
    "        logits = hidden_state.new_full((hidden_state.size(0), hidden_state.size(1), self.src_start_index+src_tokens.size(-1)),\n",
    "                                       fill_value=-1e24)\n",
    "\n",
    "        # compute eos scores\n",
    "        eos_scores = F.linear(hidden_state, self.dropout_layer(self.bart_decoder.embed_tokens.weight[2:3]))  # bsz x max_len x 1\n",
    "        \n",
    "        if self.learn_weights:   # use averge_weights compute entity labels scores\n",
    "            tag_scores = None\n",
    "            idx = 0\n",
    "            for ids in self.label_ids: # bsz x max_len x num_class\n",
    "                if len(ids) <= 1:\n",
    "                    temp_score = F.linear(hidden_state, self.dropout_layer(self.bart_decoder.embed_tokens.weight[ids]))\n",
    "                else:\n",
    "                    weight = F.softmax(self.averge_weights[idx])\n",
    "                    temp_score = F.linear(hidden_state, self.dropout_layer(self.bart_decoder.embed_tokens.weight[[ids[0]]])) * weight[0]\n",
    "                    for i in range(1, len(ids)):\n",
    "                        temp_score = temp_score + F.linear(hidden_state, self.dropout_layer(self.bart_decoder.embed_tokens.weight[[ids[i]]])) * weight[i]\n",
    "                    idx += 1\n",
    "                if tag_scores is None:\n",
    "                    tag_scores = temp_score\n",
    "                else:\n",
    "                    tag_scores = torch.cat((tag_scores, temp_score), dim=2)\n",
    "        else:\n",
    "            tag_scores = F.linear(hidden_state, self.dropout_layer(self.bart_decoder.embed_tokens.weight[self.label_start_id:self.label_end_id]))  # bsz x max_len x num_class\n",
    "\n",
    "        # bsz x max_bpe_len x hidden_size\n",
    "        src_outputs = encoder_outputs\n",
    "        if hasattr(self, 'encoder_mlp'):\n",
    "            src_outputs = self.encoder_mlp(src_outputs)\n",
    "\n",
    "        if first is not None:\n",
    "            mask = first.eq(0)  # bsz x 1 x max_word_len\n",
    "            # bsz x max_word_len x hidden_size\n",
    "            src_outputs = src_outputs.gather(index=first.unsqueeze(2).repeat(1, 1, src_outputs.size(-1)), dim=1)\n",
    "        else:\n",
    "            mask = attention_mask.eq(0)\n",
    "            # src_outputs = self.decoder.embed_tokens(src_tokens)\n",
    "        mask = mask.unsqueeze(1)\n",
    "        input_embed = self.dropout_layer(self.bart_decoder.embed_tokens(src_tokens))  # bsz x max_word_len x hidden_size\n",
    "        src_outputs = (src_outputs + input_embed)/2\n",
    "        word_scores = torch.einsum('blh,bnh->bln', hidden_state, src_outputs)  # bsz x max_len x max_word_len\n",
    "        mask = mask.__or__(src_tokens.eq(2).cumsum(dim=1).ge(1).unsqueeze(1))\n",
    "        word_scores = word_scores.masked_fill(mask, -1e32)\n",
    "\n",
    "        logits[:, :, 1:2] = eos_scores\n",
    "        logits[:, :, 2:self.src_start_index] = tag_scores\n",
    "        logits[:, :, self.src_start_index:] = word_scores\n",
    "        return logits, prompt_state\n",
    "\n",
    "    def decode(self, tokens, state):\n",
    "        return self(tokens, state)[0][:, -1]\n",
    "\n",
    "class PromptBartModel(nn.Module):\n",
    "    def __init__(self, tokenizer, label_ids, args):\n",
    "        super(PromptBartModel, self).__init__()\n",
    "        self.use_prompt = args.use_prompt\n",
    "        self.prompt_len = args.prompt_len\n",
    "        self.prompt_dim = args.prompt_dim\n",
    "        self.learn_weights = args.learn_weights\n",
    "        self.device = 'cuda' if torch.cuda.is_available else 'cpu'\n",
    "        bart_name = args.bart_name\n",
    "\n",
    "        self.bart_config = BartConfig.from_pretrained(bart_name)\n",
    "        self.bart_config.use_prompt = args.use_prompt\n",
    "        self.bart_config.preseqlen = args.prompt_len\n",
    "        bart_config = self.bart_config\n",
    "        bart_model = BartModel.from_pretrained(bart_name, config=bart_config)\n",
    "        num_tokens, _ = bart_model.encoder.embed_tokens.weight.shape\n",
    "        bart_model.resize_token_embeddings(len(tokenizer.unique_no_split_tokens)+num_tokens)\n",
    "        bart_model = avg_token_embeddings(tokenizer, bart_model, bart_name, num_tokens)\n",
    "        \n",
    "        self.prompt_encoder = PromptBartEncoder(bart_model.encoder)\n",
    "        self.prompt_decoder = PromptBartDecoder(bart_model.decoder, tokenizer.pad_token_id, label_ids, self.use_prompt, self.prompt_len, self.learn_weights)\n",
    "\n",
    "        self.prompt_inputs = torch.arange(self.prompt_len).long()\n",
    "        self.encoder_prompt_embed = nn.Embedding(self.prompt_len, bart_config.d_model)\n",
    "        self.encoder_mlp = nn.Sequential(\n",
    "                            nn.Linear(bart_config.d_model, self.prompt_dim),\n",
    "                            nn.Tanh(),\n",
    "                            nn.Linear(self.prompt_dim, bart_config.decoder_layers * 2 * bart_config.d_model))\n",
    "        \n",
    "        self.decoder_prompt_embed = nn.Embedding(self.prompt_len, bart_config.d_model)\n",
    "        self.decoder_mlp = nn.Sequential(\n",
    "                            nn.Linear(bart_config.d_model, self.prompt_dim),\n",
    "                            nn.Tanh(),\n",
    "                            nn.Linear(self.prompt_dim, bart_config.decoder_layers * 2 * bart_config.d_model))\n",
    "        \n",
    "        self.prompt_cross_embed = nn.Embedding(self.prompt_len, bart_config.d_model)\n",
    "        self.cross_mlp = nn.Sequential(\n",
    "                            nn.Linear(bart_config.d_model, self.prompt_dim),\n",
    "                            nn.Tanh(),\n",
    "                            nn.Linear(self.prompt_dim, bart_config.decoder_layers * 2 * bart_config.d_model))\n",
    "        \n",
    "        self.dropout = nn.Dropout(0.0)\n",
    "\n",
    "\n",
    "    def forward(self, src_tokens, tgt_tokens, src_seq_len, first):\n",
    "        prompt_state = self.generator(src_tokens, src_seq_len, first)\n",
    "        \n",
    "        decoder_outputs, prompt_state =  self.prompt_decoder(tgt_tokens, prompt_state)\n",
    "        return decoder_outputs\n",
    "\n",
    "\n",
    "    def generator(self, src_tokens, src_seq_len, first):\n",
    "        batch_size = src_tokens.size(0)\n",
    "        past_key_values = self.get_prompt(batch_size) if self.use_prompt else None\n",
    "        attention_mask = seq_to_mask(src_seq_len, max_len=src_tokens.size(1))\n",
    "        encoder_outputs, hidden_states = self.prompt_encoder(src_tokens, attention_mask=attention_mask, past_key_values=past_key_values)\n",
    "        prompt_state = PromptBartState(encoder_outputs, attention_mask, past_key_values, src_tokens, first, hidden_states[0], self.bart_config.preseqlen)\n",
    "\n",
    "        return prompt_state\n",
    "\n",
    " \n",
    "    def get_prompt(self, batch_size):\n",
    "        input_tokens = self.prompt_inputs.unsqueeze(0).expand(batch_size, -1).to(self.device)\n",
    "\n",
    "        # encoder prompt\n",
    "        encoder_embed = self.encoder_prompt_embed(input_tokens)\n",
    "        past_key_values = self.encoder_mlp(encoder_embed) #bsz, seqlen, layer*emb\n",
    "        bsz, seqlen, _ = past_key_values.shape\n",
    "        past_key_values = past_key_values.view(bsz, seqlen, self.bart_config.decoder_layers * 2, \n",
    "                                                self.bart_config.decoder_attention_heads, self.bart_config.d_model // self.bart_config.decoder_attention_heads)\n",
    "        past_key_values = self.dropout(past_key_values)\n",
    "        past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2) # key + value\n",
    "\n",
    "        # decoder prompt\n",
    "        decoder_embed = self.decoder_prompt_embed(input_tokens)\n",
    "        past_key_values2 = self.decoder_mlp(decoder_embed)  # bsz, seqlen, layer*emb\n",
    "        past_key_values2 = past_key_values2.view(bsz, seqlen, self.bart_config.decoder_layers * 2, \n",
    "                                                 self.bart_config.decoder_attention_heads, self.bart_config.d_model // self.bart_config.decoder_attention_heads)\n",
    "        past_key_values2 = self.dropout(past_key_values2)\n",
    "        past_key_values2 = past_key_values2.permute([2, 0, 3, 1, 4]).split(2)\n",
    "\n",
    "        # cross prompt\n",
    "        cross_embed = self.prompt_cross_embed(input_tokens)\n",
    "        past_key_values_enc = self.cross_mlp(cross_embed)  # bsz, seqlen, layer*emb\n",
    "        past_key_values_enc = past_key_values_enc.view(bsz, seqlen, self.bart_config.decoder_layers * 2, \n",
    "                                                       self.bart_config.decoder_attention_heads, self.bart_config.d_model // self.bart_config.decoder_attention_heads)\n",
    "        past_key_values_enc = self.dropout(past_key_values_enc)\n",
    "        past_key_values_enc = past_key_values_enc.permute([2, 0, 3, 1, 4]).split(2)\n",
    "\n",
    "        result = []\n",
    "        for i, key_val in enumerate(past_key_values):\n",
    "            temp_dict = {'self': {\"prev_key\": key_val[0].contiguous(),\n",
    "                                    \"prev_value\": key_val[1].contiguous(),\n",
    "                                    \"prev_key_padding_mask\": torch.zeros(bsz, seqlen).to(key_val.device).bool() #bsz, preseqlen\n",
    "                                    },\n",
    "                        }\n",
    "            key_val2 = past_key_values2[i]\n",
    "            temp_dict['encoder_decoder'] = {\"prev_key\": key_val2[0].contiguous(),\n",
    "                                            \"prev_value\": key_val2[1].contiguous(),\n",
    "                                            \"prev_key_padding_mask\": torch.zeros(bsz, seqlen).to(key_val2.device).bool()\n",
    "                                            }\n",
    "            key_val_enc = past_key_values_enc[i]\n",
    "            temp_dict['encoder'] = {\"prev_key\": key_val_enc[0].contiguous(),\n",
    "                                    \"prev_value\": key_val_enc[1].contiguous(),\n",
    "                                    \"prev_key_padding_mask\": torch.zeros(bsz, seqlen).to(key_val_enc.device).bool()\n",
    "                                    }\n",
    "            result.append(temp_dict)\n",
    "\n",
    "        return result\n",
    "\n",
    "    \n",
    "class PromptBartState(object):\n",
    "    def __init__(self, encoder_output, encoder_mask, past_key_values, src_tokens, first, src_embed_outputs, preseqlen):\n",
    "        self.encoder_output = encoder_output\n",
    "        self.encoder_mask = encoder_mask\n",
    "        self.past_key_values = past_key_values\n",
    "        self.src_tokens = src_tokens\n",
    "        self.first = first\n",
    "        self.src_embed_outputs = src_embed_outputs\n",
    "        self.preseqlen = preseqlen\n",
    "\n",
    "    def _reorder_state(self, state: Union[torch.Tensor, list, tuple], indices: torch.LongTensor, dim: int = 0):\n",
    "        if isinstance(state, torch.Tensor):\n",
    "            state = state.index_select(index=indices, dim=dim)\n",
    "        elif isinstance(state, list):\n",
    "            for i in range(len(state)):\n",
    "                assert state[i] is not None\n",
    "                state[i] = self._reorder_state(state[i], indices, dim)\n",
    "        elif isinstance(state, tuple):\n",
    "            tmp_list = []\n",
    "            for i in range(len(state)):\n",
    "                assert state[i] is not None\n",
    "                tmp_list.append(self._reorder_state(state[i], indices, dim))\n",
    "            state = tuple(tmp_list)\n",
    "        else:\n",
    "            raise TypeError(f\"Cannot reorder data of type:{type(state)}\")\n",
    "\n",
    "        return state\n",
    "\n",
    "    def reorder_state(self, indices: torch.LongTensor):\n",
    "        super().reorder_state(indices)\n",
    "        self.src_tokens = self._reorder_state(self.src_tokens, indices)\n",
    "        if self.first is not None:\n",
    "            self.first = self._reorder_state(self.first, indices)\n",
    "        self.src_embed_outputs = self._reorder_state(self.src_embed_outputs, indices)\n",
    "        if self.past_key_values is not None:\n",
    "            new = []\n",
    "            for layer in self.past_key_values:\n",
    "                new_layer = {}\n",
    "                for key1 in list(layer.keys()):\n",
    "                    new_layer_ = {}\n",
    "                    for key2 in list(layer[key1].keys()):\n",
    "                        if layer[key1][key2] is not None:\n",
    "                            layer[key1][key2] = self._reorder_state(layer[key1][key2], indices)\n",
    "                        new_layer_[key2] = layer[key1][key2]\n",
    "                    new_layer[key1] = new_layer_\n",
    "                new.append(new_layer)\n",
    "            self.past_key_values = new\n",
    "\n",
    "    def num_samples(self):\n",
    "        if self.encoder_output is not None:\n",
    "            return self.encoder_output.size(0)\n",
    "        else:\n",
    "            return None\n",
    "\n",
    "class PromptGeneratorModel(nn.Module):\n",
    "    def __init__(self, prompt_model, max_length=20, max_len_a=0.0, num_beams=1,\n",
    "                 do_sample=False, bos_token_id=None, eos_token_id=None,\n",
    "                 repetition_penalty=1, length_penalty=1.0, pad_token_id=0, restricter=None):\n",
    "        super(PromptGeneratorModel, self).__init__()\n",
    "        self.prompt_model = prompt_model\n",
    "        self.decoder = prompt_model.prompt_decoder\n",
    "\n",
    "        self.generate_func = partial(greedy_generate, decoder=self.decoder, max_length=max_length, max_len_a=max_len_a,\n",
    "                                     num_beams=num_beams,\n",
    "                                     bos_token_id=bos_token_id, eos_token_id=eos_token_id,\n",
    "                                     repetition_penalty=repetition_penalty,\n",
    "                                     length_penalty=length_penalty, pad_token_id=pad_token_id,\n",
    "                                     restricter=restricter)\n",
    "        self.do_sample = do_sample\n",
    "        self.max_length = max_length\n",
    "        self.num_beams = num_beams\n",
    "        self.bos_token_id = bos_token_id\n",
    "        self.eos_token_id = eos_token_id\n",
    "        self.repetition_penalty = repetition_penalty\n",
    "        self.length_penalty = length_penalty\n",
    "        self.pad_token_id = pad_token_id\n",
    "        self.restricter = restricter\n",
    "        self.max_len_a = max_len_a\n",
    "\n",
    "\n",
    "    def forward(self, src_tokens, tgt_tokens, src_seq_len=None, tgt_seq_len=None, first=None):\n",
    "        \"\"\"\n",
    "        :param torch.LongTensor src_tokens: bsz x max_len\n",
    "        :param torch.LongTensor tgt_tokens: bsz x max_len'\n",
    "        :param torch.LongTensor src_seq_len: bsz\n",
    "        :param torch.LongTensor tgt_seq_len: bsz\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        return self.prompt_model(src_tokens, tgt_tokens, src_seq_len, first)\n",
    "\n",
    "\n",
    "    def predict(self, src_tokens, src_seq_len=None, first=None):\n",
    "        \"\"\"\n",
    "        :param torch.LongTensor src_tokens: bsz x max_len\n",
    "        :param torch.LongTensor src_seq_len: bsz\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        prompt_state = self.prompt_model.generator(src_tokens, src_seq_len, first)    # encoder output\n",
    "        result = self.generate_func(tokens=None, state=prompt_state)\n",
    "        return result\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def greedy_generate(decoder, tokens=None, state=None, max_length=20, max_len_a=0.0, num_beams=1,\n",
    "                    bos_token_id=None, eos_token_id=None, pad_token_id=0,\n",
    "                    repetition_penalty=1, length_penalty=1.0, restricter=None):\n",
    "    if num_beams == 1:\n",
    "        token_ids = _no_beam_search_generate(decoder, tokens=tokens, state=state, max_length=max_length, max_len_a=max_len_a,\n",
    "                                             bos_token_id=bos_token_id, eos_token_id=eos_token_id,\n",
    "                                             repetition_penalty=repetition_penalty, length_penalty=length_penalty,\n",
    "                                             pad_token_id=pad_token_id, restricter=restricter)\n",
    "    else:\n",
    "        token_ids = _beam_search_generate(decoder, tokens=tokens, state=state, max_length=max_length, max_len_a=max_len_a,\n",
    "                                          num_beams=num_beams,\n",
    "                                          bos_token_id=bos_token_id, eos_token_id=eos_token_id, do_sample=False,\n",
    "                                          repetition_penalty=repetition_penalty, length_penalty=length_penalty,\n",
    "                                          pad_token_id=pad_token_id, restricter=restricter)\n",
    "\n",
    "    return token_ids\n",
    "\n",
    "\n",
    "def _no_beam_search_generate(decoder: PromptBartDecoder, state, tokens=None, max_length=20, max_len_a=0.0, bos_token_id=None,\n",
    "                             eos_token_id=None,\n",
    "                             repetition_penalty=1.0, length_penalty=1.0, pad_token_id=0,\n",
    "                             restricter=None):\n",
    "    device = get_model_device(decoder)\n",
    "    if tokens is None:\n",
    "        if bos_token_id is None:\n",
    "            raise RuntimeError(\"You have to specify either `tokens` or `bos_token_id`.\")\n",
    "        batch_size = state.num_samples()\n",
    "        if batch_size is None:\n",
    "            raise RuntimeError(\"Cannot infer the number of samples from `state`.\")\n",
    "        tokens = torch.full([batch_size, 1], fill_value=bos_token_id, dtype=torch.long).to(device)\n",
    "    batch_size = tokens.size(0)\n",
    "    if state.num_samples:\n",
    "        assert state.num_samples() == batch_size, \"The number of samples in `tokens` and `state` should match.\"\n",
    "\n",
    "    if eos_token_id is None:\n",
    "        _eos_token_id = -1\n",
    "    else:\n",
    "        _eos_token_id = eos_token_id\n",
    "    scores = decoder.decode(tokens=tokens, state=state)  # update state\n",
    "    if restricter is not None:\n",
    "        _, next_tokens = restricter(state, tokens, scores, num_beams=1)\n",
    "    else:\n",
    "        next_tokens = scores.argmax(dim=-1, keepdim=True)\n",
    "    token_ids = torch.cat([tokens, next_tokens], dim=1)\n",
    "    cur_len = token_ids.size(1)\n",
    "    dones = token_ids.new_zeros(batch_size).eq(1).__or__(next_tokens.squeeze(1).eq(eos_token_id))\n",
    "    # tokens = tokens[:, -1:]\n",
    "\n",
    "    if max_len_a!=0:\n",
    "        # (bsz x num_beams, )\n",
    "        if state.encoder_mask is not None:\n",
    "            max_lengths = (state.encoder_mask.sum(dim=1).float()*max_len_a).long() + max_length\n",
    "        else:\n",
    "            max_lengths = tokens.new_full((tokens.size(0), ), fill_value=max_length, dtype=torch.long)\n",
    "        real_max_length = max_lengths.max().item()\n",
    "    else:\n",
    "        real_max_length = max_length\n",
    "        if state.encoder_mask is not None:\n",
    "            max_lengths = state.encoder_mask.new_ones(state.encoder_mask.size(0)).long()*max_length\n",
    "        else:\n",
    "            max_lengths = tokens.new_full((tokens.size(0),), fill_value=max_length, dtype=torch.long)\n",
    "\n",
    "    while cur_len < real_max_length:\n",
    "        scores = decoder.decode(tokens=token_ids, state=state)  # batch_size x vocab_size\n",
    "\n",
    "        if repetition_penalty != 1.0:\n",
    "            token_scores = scores.gather(dim=1, index=token_ids)\n",
    "            lt_zero_mask = token_scores.lt(0).float()\n",
    "            ge_zero_mask = lt_zero_mask.eq(0).float()\n",
    "            token_scores = lt_zero_mask * repetition_penalty * token_scores + ge_zero_mask / repetition_penalty * token_scores\n",
    "            scores.scatter_(dim=1, index=token_ids, src=token_scores)\n",
    "\n",
    "        if eos_token_id is not None and length_penalty != 1.0:\n",
    "            token_scores = scores / cur_len ** length_penalty  # batch_size x vocab_size\n",
    "            eos_mask = scores.new_ones(scores.size(1))\n",
    "            eos_mask[eos_token_id] = 0\n",
    "            eos_mask = eos_mask.unsqueeze(0).eq(1)\n",
    "            scores = scores.masked_scatter(eos_mask, token_scores) \n",
    "\n",
    "        if restricter is not None:\n",
    "            _, next_tokens = restricter(state, token_ids, scores, 1)\n",
    "        else:\n",
    "            next_tokens = scores.argmax(dim=-1, keepdim=True)\n",
    "        next_tokens = next_tokens.squeeze(-1)\n",
    "\n",
    "        if _eos_token_id!=-1:\n",
    "            next_tokens = next_tokens.masked_fill(max_lengths.eq(cur_len+1), _eos_token_id)\n",
    "        next_tokens = next_tokens.masked_fill(dones, pad_token_id) \n",
    "        tokens = next_tokens.unsqueeze(1)\n",
    "\n",
    "        token_ids = torch.cat([token_ids, tokens], dim=-1)  # batch_size x max_len\n",
    "\n",
    "        end_mask = next_tokens.eq(_eos_token_id)\n",
    "        dones = dones.__or__(end_mask)\n",
    "        cur_len += 1\n",
    "\n",
    "        if dones.min() == 1:\n",
    "            break\n",
    "    return token_ids\n",
    "\n",
    "\n",
    "def _beam_search_generate(decoder: PromptBartDecoder, tokens=None, state=None, max_length=20, max_len_a=0.0, num_beams=4,\n",
    "                          bos_token_id=None, eos_token_id=None, do_sample=True,\n",
    "                          repetition_penalty=1.0, length_penalty=None, pad_token_id=0,\n",
    "                          restricter=None) -> torch.LongTensor:\n",
    "    assert do_sample is False\n",
    "    # beam search\n",
    "    device = get_model_device(decoder)\n",
    "    if tokens is None:\n",
    "        if bos_token_id is None:\n",
    "            raise RuntimeError(\"You have to specify either `tokens` or `bos_token_id`.\")\n",
    "        batch_size = state.num_samples\n",
    "        if batch_size is None:\n",
    "            raise RuntimeError(\"Cannot infer the number of samples from `state`.\")\n",
    "        tokens = torch.full([batch_size, 1], fill_value=bos_token_id, dtype=torch.long).to(device)\n",
    "    batch_size = tokens.size(0)\n",
    "    if state.num_samples:\n",
    "        assert state.num_samples == batch_size, \"The number of samples in `tokens` and `state` should match.\"\n",
    "\n",
    "    if eos_token_id is None:\n",
    "        _eos_token_id = -1\n",
    "    else:\n",
    "        _eos_token_id = eos_token_id\n",
    "\n",
    "    scores = decoder.decode(tokens=tokens, state=state)\n",
    "    vocab_size = scores.size(1)\n",
    "    assert vocab_size >= num_beams, \"num_beams should be smaller than the number of vocabulary size.\"\n",
    "\n",
    "    scores = F.log_softmax(scores, dim=-1)  # (batch_size, vocab_size)\n",
    "    if restricter is not None:\n",
    "        _next_scores, _next_tokens = restricter(state, tokens, scores, num_beams+1)\n",
    "    else:\n",
    "        # bsz x (num_beams+1)\n",
    "        _next_scores, _next_tokens = torch.topk(scores, num_beams+1, dim=1, largest=True, sorted=True)\n",
    "\n",
    "    indices = torch.arange(batch_size, dtype=torch.long).to(device)\n",
    "    indices = indices.repeat_interleave(num_beams)\n",
    "    state.reorder_state(indices)\n",
    "    tokens = tokens.index_select(dim=0, index=indices)  # batch_size * num_beams x length\n",
    "\n",
    "    if max_len_a!=0:\n",
    "        # (bsz x num_beams, )\n",
    "        if state.encoder_mask is not None:\n",
    "            max_lengths = (state.encoder_mask.sum(dim=1).float()*max_len_a).long() + max_length\n",
    "        else:\n",
    "            max_lengths = tokens.new_full((batch_size*num_beams, ), fill_value=max_length, dtype=torch.long)\n",
    "        real_max_length = max_lengths.max().item()\n",
    "    else:\n",
    "        real_max_length = max_length\n",
    "        if state.encoder_mask is not None:\n",
    "            max_lengths = state.encoder_mask.new_ones(state.encoder_mask.size(0)).long()*max_length\n",
    "        else:\n",
    "            max_lengths = tokens.new_full((batch_size*num_beams,), fill_value=max_length, dtype=torch.long)\n",
    "    hypos = [\n",
    "        BeamHypotheses(num_beams, real_max_length, length_penalty, early_stopping=False) for _ in range(batch_size)\n",
    "    ]\n",
    "\n",
    "    not_eos_mask = _next_tokens.ne(_eos_token_id)  \n",
    "    keep_mask = not_eos_mask.cumsum(dim=1).le(num_beams)  \n",
    "    keep_mask = not_eos_mask.__and__(keep_mask) \n",
    "\n",
    "    next_tokens = _next_tokens.masked_select(keep_mask).view(batch_size, num_beams) \n",
    "    next_scores = _next_scores.masked_select(keep_mask).view(batch_size, num_beams)\n",
    "\n",
    "    rows, cols = not_eos_mask.eq(0)[:, :num_beams].nonzero(as_tuple=True)\n",
    "\n",
    "    if len(rows)>0: \n",
    "        for row, col in zip(rows.tolist(), cols.tolist()):\n",
    "            _token = torch.cat([tokens[row*num_beams], _next_tokens[row, col:col+1]], dim=0)\n",
    "            hypos[row].add(_token.clone(), _next_scores[row, col].item())\n",
    "\n",
    "    # (batch_size, cur_len)\n",
    "    token_ids = torch.cat([tokens, next_tokens.view(-1, 1)], dim=-1)\n",
    "    dones = [False] * batch_size\n",
    "\n",
    "    beam_scores = next_scores.view(-1)  # batch_size * num_beams\n",
    "\n",
    "    cur_len = token_ids.size(1)\n",
    "\n",
    "    # 0, num_beams, 2*num_beams, ...\n",
    "    batch_inds_with_numbeams_interval = (torch.arange(batch_size) * num_beams).view(-1, 1).to(token_ids)\n",
    "\n",
    "    while cur_len < real_max_length:\n",
    "        scores = decoder.decode(token_ids, state)  # (bsz x num_beams, vocab_size)\n",
    "        if repetition_penalty != 1.0:\n",
    "            token_scores = scores.gather(dim=1, index=token_ids)\n",
    "            lt_zero_mask = token_scores.lt(0).float()\n",
    "            ge_zero_mask = lt_zero_mask.eq(0).float()\n",
    "            token_scores = lt_zero_mask * repetition_penalty * token_scores + ge_zero_mask / repetition_penalty * token_scores\n",
    "            scores.scatter_(dim=1, index=token_ids, src=token_scores)\n",
    "\n",
    "        if _eos_token_id!=-1:\n",
    "            max_len_eos_mask = max_lengths.eq(cur_len+1)\n",
    "            eos_scores = scores[:, _eos_token_id]\n",
    "            scores[:, _eos_token_id] = torch.where(max_len_eos_mask, eos_scores+1e32, eos_scores)\n",
    "\n",
    "        scores = F.log_softmax(scores, dim=-1)  # (batch_size * num_beams, vocab_size)\n",
    "        _scores = scores + beam_scores[:, None]  # (batch_size * num_beams, vocab_size)\n",
    "        _scores = _scores.view(batch_size, -1)  # (batch_size, num_beams*vocab_size)\n",
    " \n",
    "        if restricter is not None:\n",
    "            next_scores, ids = restricter(state, token_ids, _scores, 2 * num_beams)\n",
    "        else:\n",
    "            next_scores, ids = torch.topk(_scores, 2 * num_beams, dim=1, largest=True, sorted=True)  # (bsz, 2*num_beams)\n",
    "        from_which_beam = ids // vocab_size  # (batch_size, 2*num_beams)\n",
    "        next_tokens = ids % vocab_size  # (batch_size, 2*num_beams)\n",
    "\n",
    "        not_eos_mask = next_tokens.ne(_eos_token_id)  \n",
    "        keep_mask = not_eos_mask.cumsum(dim=1).le(num_beams)  \n",
    "        keep_mask = not_eos_mask.__and__(keep_mask) \n",
    "\n",
    "        _next_tokens = next_tokens.masked_select(keep_mask).view(-1, 1)\n",
    "        _from_which_beam = from_which_beam.masked_select(keep_mask).view(batch_size, num_beams) \n",
    "        _next_scores = next_scores.masked_select(keep_mask).view(batch_size, num_beams)\n",
    "        beam_scores = _next_scores.view(-1)\n",
    "\n",
    "        flag = True\n",
    "        if cur_len+1 == real_max_length:\n",
    "            eos_batch_idx = torch.arange(batch_size).to(next_tokens).repeat_interleave(repeats=num_beams, dim=0)\n",
    "            eos_beam_ind = torch.arange(num_beams).to(token_ids).repeat(batch_size) \n",
    "            eos_beam_idx = from_which_beam[:, :num_beams].reshape(-1) \n",
    "        else:\n",
    "            effective_eos_mask = next_tokens[:, :num_beams].eq(_eos_token_id)  # batch_size x num_beams\n",
    "            if effective_eos_mask.sum().gt(0):\n",
    "                eos_batch_idx, eos_beam_ind = effective_eos_mask.nonzero(as_tuple=True)\n",
    "                eos_beam_idx = eos_batch_idx * num_beams * 2 + eos_beam_ind\n",
    "                eos_beam_idx = from_which_beam.view(-1)[eos_beam_idx]\n",
    "            else:\n",
    "                flag = False\n",
    "\n",
    "        if flag:\n",
    "            _token_ids = torch.cat([token_ids, _next_tokens], dim=-1)\n",
    "            for batch_idx, beam_ind, beam_idx in zip(eos_batch_idx.tolist(), eos_beam_ind.tolist(),\n",
    "                                                     eos_beam_idx.tolist()):\n",
    "                if not dones[batch_idx]:\n",
    "                    score = next_scores[batch_idx, beam_ind].item()\n",
    "                    if _eos_token_id!=-1:\n",
    "                        hypos[batch_idx].add(_token_ids[batch_idx * num_beams + beam_idx, :cur_len].clone(), score)\n",
    "                    else:\n",
    "                        hypos[batch_idx].add(_token_ids[batch_idx * num_beams + beam_idx].clone(), score)\n",
    "\n",
    "        reorder_inds = (batch_inds_with_numbeams_interval + _from_which_beam).view(-1)  # flatten\n",
    "        state.reorder_state(reorder_inds)\n",
    "        token_ids = torch.cat([token_ids.index_select(index=reorder_inds, dim=0), _next_tokens], dim=-1)\n",
    "\n",
    "        for batch_idx in range(batch_size):\n",
    "            dones[batch_idx] = dones[batch_idx] or hypos[batch_idx].is_done(next_scores[batch_idx, 0].item()) or \\\n",
    "                               max_lengths[batch_idx*num_beams]==cur_len+1\n",
    "\n",
    "        cur_len += 1\n",
    "\n",
    "        if all(dones):\n",
    "            break\n",
    "\n",
    "    # select the best hypotheses\n",
    "    tgt_len = token_ids.new_zeros(batch_size)\n",
    "    best = []\n",
    "\n",
    "    for i, hypotheses in enumerate(hypos):\n",
    "        best_hyp = max(hypotheses.hyp, key=lambda x: x[0])[1]\n",
    "        if _eos_token_id!=-1:\n",
    "            best_hyp = torch.cat([best_hyp, best_hyp.new_ones(1)*_eos_token_id])\n",
    "        tgt_len[i] = len(best_hyp)\n",
    "        best.append(best_hyp)\n",
    "\n",
    "    # generate target batch\n",
    "    decoded = token_ids.new_zeros(batch_size, tgt_len.max().item()).fill_(pad_token_id)\n",
    "    for i, hypo in enumerate(best):\n",
    "        decoded[i, :tgt_len[i]] = hypo\n",
    "\n",
    "    return decoded\n",
    "\n",
    "\n",
    "class BeamHypotheses(object):\n",
    "    def __init__(self, num_beams, max_length, length_penalty, early_stopping):\n",
    "        \"\"\"\n",
    "        Initialize n-best list of hypotheses.\n",
    "        \"\"\"\n",
    "        self.max_length = max_length - 1  # ignoring bos_token\n",
    "        self.length_penalty = length_penalty\n",
    "        self.early_stopping = early_stopping\n",
    "        self.num_beams = num_beams\n",
    "        self.hyp = []\n",
    "        self.worst_score = 1e9\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"\n",
    "        Number of hypotheses in the list.\n",
    "        \"\"\"\n",
    "        return len(self.hyp)\n",
    "\n",
    "    def add(self, hyp, sum_logprobs):\n",
    "        \"\"\"\n",
    "        Add a new hypothesis to the list.\n",
    "        \"\"\"\n",
    "        score = sum_logprobs / len(hyp) ** self.length_penalty\n",
    "        if len(self) < self.num_beams or score > self.worst_score:\n",
    "            self.hyp.append((score, hyp))\n",
    "            if len(self) > self.num_beams:\n",
    "                sorted_scores = sorted([(s, idx) for idx, (s, _) in enumerate(self.hyp)])\n",
    "                del self.hyp[sorted_scores[0][1]]\n",
    "                self.worst_score = sorted_scores[1][0]\n",
    "            else:\n",
    "                self.worst_score = min(score, self.worst_score)\n",
    "\n",
    "    def is_done(self, best_sum_logprobs):\n",
    "        \"\"\"\n",
    "        If there are enough hypotheses and that none of the hypotheses being generated\n",
    "        can become better than the worst one in the heap, then we are done with this sentence.\n",
    "        \"\"\"\n",
    "        if len(self) < self.num_beams:\n",
    "            return False\n",
    "        elif self.early_stopping:\n",
    "            return True\n",
    "        else:\n",
    "            return self.worst_score >= best_sum_logprobs / self.max_length ** self.length_penalty\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b33fae4",
   "metadata": {},
   "source": [
    "## Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27f3af06",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import optim\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "\n",
    "class Trainer(object):\n",
    "    def __init__(self, train_data=None, dev_data=None, test_data=None, model=None, process=None, args=None, logger=None, loss=None, metrics=None, writer=None) -> None:\n",
    "        self.train_data = train_data\n",
    "        self.dev_data = dev_data\n",
    "        self.test_data = test_data\n",
    "        self.model = model\n",
    "        self.process = process\n",
    "        self.logger = logger\n",
    "        self.metrics = metrics\n",
    "        self.writer = writer\n",
    "        self.loss = loss\n",
    "        self.num_epochs = args.num_epochs\n",
    "        self.batch_size = args.batch_size\n",
    "        self.lr = args.learning_rate\n",
    "        self.eval_begin_epoch = args.eval_begin_epoch\n",
    "        self.device = args.device\n",
    "        self.load_path = args.load_path\n",
    "        self.save_path = args.save_path\n",
    "        self.refresh_step = 1\n",
    "        self.best_metric = 0\n",
    "        self.best_dev_epoch = None\n",
    "        self.optimizer = None\n",
    "        if self.train_data is not None:\n",
    "            self.train_num_steps = len(self.train_data) * self.num_epochs\n",
    "        self.step = 0\n",
    "        self.args = args\n",
    "        \n",
    "    \n",
    "    def train(self):\n",
    "        self.before_train()  # something should do before training\n",
    "        self.step = 0\n",
    "        self.model.train()\n",
    "        self.logger.info(\"***** Running training *****\")\n",
    "        self.logger.info(\"  Num instance = %d\", len(self.train_data)*self.batch_size)\n",
    "        self.logger.info(\"  Num epoch = %d\", self.num_epochs)\n",
    "        self.logger.info(\"  Batch size = %d\", self.batch_size)\n",
    "        self.logger.info(\"  Learning rate = {}\".format(self.lr))\n",
    "        self.logger.info(\"  Evaluate begin = %d\", self.eval_begin_epoch)\n",
    "\n",
    "        if self.load_path is not None:  # load model from load_path\n",
    "            self.logger.info(\"Loading model from {}\".format(self.load_path))\n",
    "            load_model_dict = torch.load(self.args.load_path)\n",
    "            model_dict = self.model.state_dict()\n",
    "            for name in load_model_dict:\n",
    "                if name in model_dict:\n",
    "                    if model_dict[name].shape == load_model_dict[name].shape:\n",
    "                        model_dict[name] = load_model_dict[name]\n",
    "                    else:\n",
    "                        self.logger.info(f\"Skip loading parameter: {name}, \"\n",
    "                            f\"required shape: {model_dict[name].shape}, \"\n",
    "                            f\"loaded shape: {load_model_dict[name].shape}\")\n",
    "                else:\n",
    "                    self.logger.info(f\"Not Found! Skip loading parameter: {name}.\")\n",
    "            self.model.load_state_dict(model_dict)\n",
    "            self.logger.info(\"Load model successful!\")\n",
    "\n",
    "        with tqdm(total=self.train_num_steps, postfix='loss:{0:<6.5f}', leave=False, dynamic_ncols=True, initial=self.step) as pbar:\n",
    "            self.pbar = pbar\n",
    "            avg_loss = 0\n",
    "            for epoch in range(self.num_epochs):\n",
    "                pbar.set_description_str(desc=\"Epoch {}/{}\".format(epoch, self.num_epochs))\n",
    "                for batch in self.train_data:\n",
    "                    self.step += 1\n",
    "                    batch = (tup.to(self.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)\n",
    "                    loss = self._step(batch, mode=\"train\")\n",
    "                    avg_loss += loss.item()\n",
    "\n",
    "                    self.optimizer.zero_grad()\n",
    "                    loss.backward()\n",
    "                    self.optimizer.step()\n",
    "\n",
    "                    if self.step % self.refresh_step == 0:\n",
    "                        avg_loss = float(avg_loss) / self.refresh_step\n",
    "                        print_output = \"loss:{:<6.5f}\".format(avg_loss)\n",
    "                        pbar.update(1)\n",
    "                        pbar.set_postfix_str(print_output)\n",
    "                        self.writer.add_scalar(tag='loss', scalar_value=avg_loss, global_step=self.step)    # tensorbordx\n",
    "                        avg_loss = 0\n",
    "                if epoch >= self.eval_begin_epoch:\n",
    "                    self.evaluate(epoch)   # generator to dev.\n",
    "            pbar.close()\n",
    "            self.pbar = None\n",
    "            self.logger.info(\"Get best performance at epoch {}, best f1 score is {:.2f}\".format(self.best_dev_epoch, self.best_metric))\n",
    "\n",
    "\n",
    "\n",
    "    def evaluate(self, epoch):\n",
    "        self.model.eval()\n",
    "        self.logger.info(\"***** Running evaluate *****\")\n",
    "        self.logger.info(\"  Num instance = %d\", len(self.dev_data)*self.batch_size)\n",
    "        self.logger.info(\"  Batch size = %d\", self.batch_size)\n",
    "        with torch.no_grad():\n",
    "            with tqdm(total=len(self.dev_data), leave=False, dynamic_ncols=True) as pbar:\n",
    "                pbar.set_description_str(desc=\"Dev\")\n",
    "                for batch in self.dev_data:\n",
    "                    batch = (tup.to(self.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)  # to cpu/cuda device\n",
    "                    self._step(batch, mode=\"dev\")\n",
    "                    pbar.update()\n",
    "                # evaluate done\n",
    "                eva_result = self.metrics.get_metric()\n",
    "                pbar.close()\n",
    "                self.logger.info(\"Epoch {}/{}, best f1: {}, current f1 score: {:.2f}, recall: {:.2f}, precision: {:.2f}.\"\\\n",
    "                            .format(epoch, self.num_epochs, self.best_metric, eva_result['f'], eva_result['rec'], eva_result['pre']))\n",
    "                self.writer.add_scalars('evaluate', {'f1': eva_result['f'],\n",
    "                                                    'recall': eva_result['rec'],\n",
    "                                                    'precision': eva_result['pre']}, epoch)\n",
    "                if eva_result['f'] >= self.best_metric:  # this epoch get best performance\n",
    "                    self.logger.info(\"Get better performance at epoch {}\".format(epoch))\n",
    "                    self.best_dev_epoch = epoch\n",
    "                    self.best_metric = eva_result['f']  # update best metric(f1 score)\n",
    "                    if self.save_path is not None:      # need to save model\n",
    "                        torch.save(self.model.state_dict(), self.save_path+\"/best_model.pth\")\n",
    "                        self.logger.info(\"Save best model at {}\".format(self.save_path))\n",
    "\n",
    "        self.model.train()\n",
    "\n",
    "    def predict(self):\n",
    "        self.model.eval()\n",
    "        self.logger.info(\"***** Running testing *****\")\n",
    "        self.logger.info(\"  Num instance = %d\", len(self.test_data)*self.batch_size)\n",
    "        self.logger.info(\"  Batch size = %d\", self.batch_size)\n",
    "        if self.load_path is not None:  # load model from load_path\n",
    "            self.logger.info(\"Loading model from {}\".format(self.load_path))\n",
    "            self.model.load_state_dict(torch.load(self.load_path))\n",
    "            self.logger.info(\"Load model successful!\")\n",
    "            self.model.to(self.device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            with tqdm(total=len(self.test_data), leave=False, dynamic_ncols=True) as pbar:\n",
    "                pbar.set_description_str(desc=\"Test\")\n",
    "                texts = []\n",
    "                labels = []\n",
    "                for batch in self.test_data:\n",
    "                    batch = (tup.to(self.device)  if isinstance(tup, torch.Tensor) else tup for tup in batch)  # to cpu/cuda device\n",
    "                    src_tokens, src_seq_len, first, raw_words = batch\n",
    "                    preds = self._step((src_tokens, src_seq_len, first), mode=\"test\")\n",
    "                    outputs = convert_preds_to_outputs(preds, raw_words, self.process.mapping, self.process.tokenizer)\n",
    "                    texts.extend(raw_words)\n",
    "                    labels.extend(outputs)\n",
    "                    pbar.update()\n",
    "\n",
    "        self.logger.info(\"***** Predict example *****\")\n",
    "        idx = random.randint(0, len(texts))\n",
    "        print(len(texts), len(labels))\n",
    "        self.logger.info(\"Raw texts: \" + \" \".join(texts[idx]))\n",
    "        self.logger.info(\"Prediction: \" + \" \".join(labels[idx]))\n",
    "        if self.args.write_path is not None:    # write predict\n",
    "            write_predictions(self.args.write_path, texts, labels)\n",
    "            self.logger.info(\"Write into {}!\".format(self.args.write_path))\n",
    "\n",
    "        \n",
    "\n",
    "    def _step(self, batch, mode=\"train\"):\n",
    "        if mode==\"dev\": # dev: compute metric\n",
    "            src_tokens, tgt_tokens, src_seq_len, tgt_seq_len, first, target_span = batch\n",
    "            pred = self.model.predict(src_tokens, src_seq_len, first)\n",
    "            self.metrics.evaluate(target_span, pred, tgt_tokens)\n",
    "            return\n",
    "        elif mode==\"test\":  # test: just get pred\n",
    "            src_tokens, src_seq_len, first = batch\n",
    "            pred = self.model.predict(src_tokens, src_seq_len, first)\n",
    "            return pred\n",
    "        else:   # train: get loss\n",
    "            src_tokens, tgt_tokens, src_seq_len, tgt_seq_len, first, target_span = batch\n",
    "            pred = self.model(src_tokens, tgt_tokens, src_seq_len, first)\n",
    "            loss = self.loss(tgt_tokens, tgt_seq_len, pred)\n",
    "            return loss\n",
    "\n",
    "\n",
    "    def before_train(self):\n",
    "        parameters = []\n",
    "        params = {'lr':self.lr, 'weight_decay':1e-2}\n",
    "        params['params'] = [param for name, param in self.model.named_parameters() if not ('bart_encoder' in name or 'bart_decoder' in name)]\n",
    "        parameters.append(params)\n",
    "\n",
    "        params = {'lr':self.lr, 'weight_decay':1e-2}\n",
    "        params['params'] = []\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if ('bart_encoder' in name or 'bart_decoder' in name) and not ('layernorm' in name or 'layer_norm' in name):\n",
    "                params['params'].append(param)\n",
    "        parameters.append(params)\n",
    "\n",
    "        params = {'lr':self.lr, 'weight_decay':0}\n",
    "        params['params'] = []\n",
    "        for name, param in self.model.named_parameters():\n",
    "            if ('bart_encoder' in name or 'bart_decoder' in name) and ('layernorm' in name or 'layer_norm' in name):\n",
    "                params['params'].append(param)\n",
    "        parameters.append(params)\n",
    "\n",
    "        self.optimizer = optim.AdamW(parameters)\n",
    "\n",
    "        if self.args.freeze_plm:    # freeze pretrained language model(bart)\n",
    "            for name, par in self.model.named_parameters():\n",
    "                if 'prompt_encoder' in name or 'prompt_decoder' in name and \"bart_mlp\" not in name:\n",
    "                    par.requires_grad = False\n",
    "\n",
    "        self.model.to(self.device)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edc9566d",
   "metadata": {},
   "source": [
    "## Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34de8e5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "class Seq2SeqSpanMetric(object):\n",
    "    def __init__(self, eos_token_id, num_labels, target_type='word'):\n",
    "        self.eos_token_id = eos_token_id\n",
    "        self.num_labels = num_labels\n",
    "        self.word_start_index = num_labels+2\n",
    "\n",
    "        self.fp = 0\n",
    "        self.tp = 0\n",
    "        self.fn = 0\n",
    "        self.em = 0\n",
    "        self.total = 0\n",
    "        self.target_type = target_type \n",
    "\n",
    "    def evaluate(self, target_span, pred, tgt_tokens):\n",
    "        self.total += pred.size(0)\n",
    "        pred_eos_index = pred.flip(dims=[1]).eq(self.eos_token_id).cumsum(dim=1).long()\n",
    "        target_eos_index = tgt_tokens.flip(dims=[1]).eq(self.eos_token_id).cumsum(dim=1).long()\n",
    "\n",
    "        pred = pred[:, 1:]  \n",
    "        tgt_tokens = tgt_tokens[:, 1:]\n",
    "        pred_seq_len = pred_eos_index.flip(dims=[1]).eq(pred_eos_index[:, -1:]).sum(dim=1) # bsz\n",
    "        pred_seq_len = (pred_seq_len - 2).tolist()\n",
    "        target_seq_len = target_eos_index.flip(dims=[1]).eq(target_eos_index[:, -1:]).sum(dim=1) # bsz\n",
    "        target_seq_len = (target_seq_len-2).tolist()\n",
    "        pred_spans = []\n",
    "        for i, (ts, ps) in enumerate(zip(target_span, pred.tolist())):\n",
    "            em = 0\n",
    "            ps = ps[:pred_seq_len[i]]\n",
    "            if pred_seq_len[i]==target_seq_len[i]:\n",
    "                em = int(tgt_tokens[i, :target_seq_len[i]].eq(pred[i, :target_seq_len[i]]).sum().item()==target_seq_len[i])\n",
    "            self.em += em\n",
    "            pairs = []\n",
    "            cur_pair = []\n",
    "            if len(ps):\n",
    "                for j in ps:\n",
    "                    if j<self.word_start_index:\n",
    "                        if self.target_type == 'span':\n",
    "                            if len(cur_pair)>0 and len(cur_pair)%2==0:\n",
    "                                if all([cur_pair[i]<=cur_pair[i+1] for i in range(len(cur_pair)-1)]):\n",
    "                                    pairs.append(tuple(cur_pair+[j]))\n",
    "                        else:\n",
    "                            if len(cur_pair) > 0:\n",
    "                                if all([cur_pair[i]<cur_pair[i+1] for i in range(len(cur_pair)-1)]):\n",
    "                                    pairs.append(tuple(cur_pair + [j]))\n",
    "                        cur_pair = []\n",
    "                    else:\n",
    "                        cur_pair.append(j)\n",
    "            pred_spans.append(pairs.copy())\n",
    "\n",
    "            tp, fn, fp = _compute_tp_fn_fp(pairs, ts)\n",
    "            self.fn += fn\n",
    "            self.tp += tp\n",
    "            self.fp += fp\n",
    "\n",
    "    def get_metric(self, reset=True):\n",
    "        res = {}\n",
    "        f, pre, rec = _compute_f_pre_rec(1, self.tp, self.fn, self.fp)\n",
    "        res['f'] = round(f, 4)*100\n",
    "        res['rec'] = round(rec, 4)*100\n",
    "        res['pre'] = round(pre, 4)*100\n",
    "        res['em'] = round(self.em/self.total, 4)\n",
    "        if reset:\n",
    "            self.total = 0\n",
    "            self.fp = 0\n",
    "            self.tp = 0\n",
    "            self.fn = 0\n",
    "            self.em = 0\n",
    "        return res\n",
    "\n",
    "def _compute_f_pre_rec(beta_square, tp, fn, fp):\n",
    "    r\"\"\"\n",
    "\n",
    "    :param tp: int, true positive\n",
    "    :param fn: int, false negative\n",
    "    :param fp: int, false positive\n",
    "    :return: (f, pre, rec)\n",
    "    \"\"\"\n",
    "    pre = tp / (fp + tp + 1e-13)\n",
    "    rec = tp / (fn + tp + 1e-13)\n",
    "    f = (1 + beta_square) * pre * rec / (beta_square * pre + rec + 1e-13)\n",
    "\n",
    "    return f, pre, rec\n",
    "\n",
    "\n",
    "def _compute_tp_fn_fp(ps, ts):\n",
    "    ps = ps.copy()\n",
    "    tp = 0\n",
    "    fp = 0\n",
    "    fn = 0\n",
    "    if isinstance(ts, (set, list, np.ndarray)):\n",
    "        ts = {tuple(key):1 for key in list(ts)}\n",
    "    if isinstance(ps, (set, list, np.ndarray)):\n",
    "        ps = {tuple(key):1 for key in list(ps)}\n",
    "\n",
    "    for key in ts.keys():\n",
    "        t_num = ts[key]\n",
    "        if key not in ps:\n",
    "            p_num = 0\n",
    "        else:\n",
    "            p_num = ps[key]\n",
    "        tp += min(p_num, t_num)\n",
    "        fp += max(p_num - t_num, 0)\n",
    "        fn += max(t_num - p_num, 0)\n",
    "        if key in ps:\n",
    "            ps.pop(key)\n",
    "    fp += sum(ps.values())\n",
    "    return tp, fn, fp\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16b09372",
   "metadata": {},
   "source": [
    "## Fine-Tune\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61c96cfa",
   "metadata": {},
   "source": [
    "### Train on CoNLL2003"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "835efde0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration file of model parameters\n",
    "class Config(object):\n",
    "    dataset_name = \"conll2003\"\n",
    "    bart_name = \"facebook/bart-large\"\n",
    "    num_epochs = 30\n",
    "    device = \"cuda\"\n",
    "    batch_size = 16\n",
    "    learning_rate = 2e-5\n",
    "    warmup_ratio = 0.01\n",
    "    eval_begin_epoch = 16\n",
    "    num_beams = 1\n",
    "    seed = 1\n",
    "    length_penalty = 1\n",
    "    tgt_max_len = 10\n",
    "    src_seq_ratio = 0.6\n",
    "    prompt_len = 10\n",
    "    prompt_dim = 800\n",
    "    load_path = None\n",
    "    save_path = None    # save model\n",
    "    write_path = None\n",
    "    notes = \"\"\n",
    "    use_prompt = True\n",
    "    learn_weights = True\n",
    "    freeze_plm = True   # freeze bart model\n",
    "    do_train = True\n",
    "    do_test = False\n",
    "\n",
    "    \n",
    "cfg = Config()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ca6614",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_CLASS = {\n",
    "    'conll2003': ConllNERDataset,\n",
    "    'mit-movie': ConllNERDataset,\n",
    "    'mit-restaurant': ConllNERDataset,\n",
    "    'atis': ConllNERDataset\n",
    "}\n",
    "\n",
    "DATA_PROCESS = {\n",
    "    'conll2003': ConllNERProcessor,\n",
    "    'mit-movie': ConllNERProcessor,\n",
    "    'mit-restaurant': ConllNERProcessor,\n",
    "    'atis': ConllNERProcessor\n",
    "}\n",
    "\n",
    "DATA_PATH = {\n",
    "    'conll2003': {'train': 'data/conll2003/train.txt',\n",
    "                  'dev': 'data/conll2003/dev.txt',\n",
    "                  'test': 'data/conll2003/test.txt'},\n",
    "    'mit-movie': {'train': 'data/mit-movie/20-shot-train.txt',\n",
    "                  'dev': 'data/mit-movie/test.txt'},\n",
    "    'mit-restaurant': {'train': 'data/mit-restaurant/10-shot-train.txt',\n",
    "                  'dev': 'data/mit-restaurant/test.txt'},\n",
    "    'atis': {'train': 'data/atis/20-shot-train.txt',\n",
    "                  'dev': 'data/atis/test.txt'}\n",
    "}\n",
    "\n",
    "MAPPING = {\n",
    "    'conll2003': {'loc': '<<location>>',\n",
    "                'per': '<<person>>',\n",
    "                'org': '<<organization>>',\n",
    "                'misc': '<<others>>'},\n",
    "    'mit-movie': mit_movie_mapping,\n",
    "    'mit-restaurant': mit_restaurant_mapping,\n",
    "    'atis': atis_mapping\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "333004d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose data_path, data_class, data_mapping and set seed\n",
    "data_path = DATA_PATH[cfg.dataset_name]\n",
    "dataset_class, data_process = DATASET_CLASS[cfg.dataset_name], DATA_PROCESS[cfg.dataset_name]\n",
    "mapping = MAPPING[cfg.dataset_name]\n",
    "\n",
    "set_seed(cfg.seed) # set seed, default is 1\n",
    "if cfg.save_path is not None:  # make save_path dir\n",
    "    cfg.save_path = os.path.join(cfg.save_path, cfg.dataset_name+\"_\"+str(cfg.batch_size)+\"_\"+str(cfg.learning_rate)+cfg.notes)\n",
    "    if not os.path.exists(cfg.save_path):\n",
    "        os.makedirs(cfg.save_path, exist_ok=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5da9e180",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train\n",
    "if cfg.do_train:\n",
    "    process = data_process(data_path=data_path, mapping=mapping, bart_name=cfg.bart_name, learn_weights=cfg.learn_weights)\n",
    "    train_dataset = dataset_class(data_processor=process, mode='train')\n",
    "    train_dataloader = DataLoader(train_dataset, collate_fn=train_dataset.collate_fn, batch_size=cfg.batch_size, num_workers=4)\n",
    "    \n",
    "    dev_dataset = dataset_class(data_processor=process, mode='dev')\n",
    "    dev_dataloader = DataLoader(dev_dataset, collate_fn=dev_dataset.collate_fn, batch_size=cfg.batch_size, num_workers=4)\n",
    "\n",
    "    label_ids = list(process.mapping2id.values())\n",
    "\n",
    "    prompt_model = PromptBartModel(tokenizer=process.tokenizer, label_ids=label_ids, args=cfg)\n",
    "    model = PromptGeneratorModel(prompt_model=prompt_model, bos_token_id=0,\n",
    "                                eos_token_id=1,\n",
    "                                max_length=cfg.tgt_max_len, max_len_a=cfg.src_seq_ratio,num_beams=cfg.num_beams, do_sample=False,\n",
    "                                repetition_penalty=1, length_penalty=cfg.length_penalty, pad_token_id=1,\n",
    "                                restricter=None)\n",
    "    metrics = Seq2SeqSpanMetric(eos_token_id=1, num_labels=len(label_ids), target_type='word')\n",
    "    loss = get_loss\n",
    "\n",
    "    trainer = Trainer(train_data=train_dataloader, dev_data=dev_dataloader, test_data=None, model=model, args=cfg, logger=logger, loss=loss, metrics=metrics, writer=writer)\n",
    "    trainer.train()\n",
    "\n",
    "    writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "937ffb04",
   "metadata": {},
   "source": [
    "### Train on few-shot datasets\n",
    "Set config.dataset = \"few-shot dataset\" (mit-movie, mit-restaurant, atis), and change lr, batch_Size.\n",
    "\n",
    "If don't load model dict, set config.load_path = None, else set config.load_path = \"ckpt_path\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68b6e470",
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg.learning_rate = 1e-4\n",
    "cfg.src_seq_ratio = 0.8\n",
    "cfg.dataset_name = \"mit-movie\"\n",
    "cfg.load_path = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5827004",
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose data_path, data_class, data_mapping and set seed\n",
    "data_path = DATA_PATH[cfg.dataset_name]\n",
    "dataset_class, data_process = DATASET_CLASS[cfg.dataset_name], DATA_PROCESS[cfg.dataset_name]\n",
    "mapping = MAPPING[cfg.dataset_name]\n",
    "\n",
    "set_seed(cfg.seed) # set seed, default is 1\n",
    "if cfg.save_path is not None:  # make save_path dir\n",
    "    cfg.save_path = os.path.join(cfg.save_path, cfg.dataset_name+\"_\"+str(cfg.batch_size)+\"_\"+str(cfg.learning_rate)+cfg.notes)\n",
    "    if not os.path.exists(cfg.save_path):\n",
    "        os.makedirs(cfg.save_path, exist_ok=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55b5c3d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train\n",
    "if cfg.do_train:\n",
    "    process = data_process(data_path=data_path, mapping=mapping, bart_name=cfg.bart_name, learn_weights=cfg.learn_weights)\n",
    "    train_dataset = dataset_class(data_processor=process, mode='train')\n",
    "    train_dataloader = DataLoader(train_dataset, collate_fn=train_dataset.collate_fn, batch_size=cfg.batch_size, num_workers=4)\n",
    "    \n",
    "    dev_dataset = dataset_class(data_processor=process, mode='dev')\n",
    "    dev_dataloader = DataLoader(dev_dataset, collate_fn=dev_dataset.collate_fn, batch_size=cfg.batch_size, num_workers=4)\n",
    "\n",
    "    label_ids = list(process.mapping2id.values())\n",
    "\n",
    "    prompt_model = PromptBartModel(tokenizer=process.tokenizer, label_ids=label_ids, args=cfg)\n",
    "    model = PromptGeneratorModel(prompt_model=prompt_model, bos_token_id=0,\n",
    "                                eos_token_id=1,\n",
    "                                max_length=cfg.tgt_max_len, max_len_a=cfg.src_seq_ratio,num_beams=cfg.num_beams, do_sample=False,\n",
    "                                repetition_penalty=1, length_penalty=cfg.length_penalty, pad_token_id=1,\n",
    "                                restricter=None)\n",
    "    metrics = Seq2SeqSpanMetric(eos_token_id=1, num_labels=len(label_ids), target_type='word')\n",
    "    loss = get_loss\n",
    "\n",
    "    trainer = Trainer(train_data=train_dataloader, dev_data=dev_dataloader, test_data=None, model=model, args=cfg, logger=logger, loss=loss, metrics=metrics, writer=writer)\n",
    "    trainer.train()\n",
    "\n",
    "    writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b84bcf6d",
   "metadata": {},
   "source": [
    "### Evaluate\n",
    "Set eval_begin_epoch, and evaluate when epoch is greater than or equal to eval_begin_epoch during training.\n",
    "\n",
    "You can also get trainer, and call trainer.evaluate(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "740a2917",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd16b4a4",
   "metadata": {},
   "source": [
    "## Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a74d8c1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration file of model parameters\n",
    "cfg.dataset_name = \"mit-movie\"\n",
    "cfg.load_path = \"save_models/conll2003_16_2e-05/best_model.pth\"\n",
    "cfg.write_path = \"prediction.txt\"\n",
    "cfg.save_path = None\n",
    "cfg.do_train = False\n",
    "cfg.do_test = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8eafd135",
   "metadata": {},
   "outputs": [],
   "source": [
    "process = data_process(data_path=data_path, mapping=mapping, bart_name=cfg.bart_name, learn_weights=cfg.learn_weights)\n",
    "test_dataset = dataset_class(data_processor=process, mode='test')\n",
    "test_dataloader = DataLoader(test_dataset, collate_fn=test_dataset.collate_fn, batch_size=cfg.batch_size, num_workers=4)\n",
    "\n",
    "label_ids = list(process.mapping2id.values())\n",
    "prompt_model = PromptBartModel(tokenizer=process.tokenizer, label_ids=label_ids, args=cfg)\n",
    "model = PromptGeneratorModel(prompt_model=prompt_model, bos_token_id=0,\n",
    "                            eos_token_id=1,\n",
    "                            max_length=cfg.tgt_max_len, max_len_a=cfg.src_seq_ratio,num_beams=cfg.num_beams, do_sample=False,\n",
    "                            repetition_penalty=1, length_penalty=cfg.length_penalty, pad_token_id=1,\n",
    "                            restricter=None)\n",
    "trainer = Trainer(train_data=None, dev_data=None, test_data=test_dataloader, model=model, process=process, args=cfg, logger=logger, loss=None, metrics=None, writer=None)\n",
    "trainer.predict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86ab2c86",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "d169b6375943c11c98a1b56380b0356505937e80ef2527489143fd65051338ee"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
